import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

head_counter=[0 for i in range(10)]    # 查看专家分布是否均匀

#网络定义
class moe_net(nn.Module):
    def __init__(self):
        super(moe_net,self).__init__()
        # 256*256
        self.conv1=nn.Conv2d(in_channels=3, out_channels=6,kernel_size=3,stride=1,padding=1)
        self.pool1=nn.MaxPool2d(kernel_size=2, stride=2)
        # 128*128
        self.conv2=nn.Conv2d(in_channels=6, out_channels=12,kernel_size=3,stride=1,padding=1)
        self.pool2=nn.MaxPool2d(kernel_size=2, stride=2)
        # 64*64
        self.conv3=nn.Conv2d(in_channels=12, out_channels=24,kernel_size=3,stride=1,padding=1)
        self.pool3=nn.MaxPool2d(kernel_size=2, stride=2)
        # 32*32
        self.conv4=nn.Conv2d(in_channels=24, out_channels=48,kernel_size=3,stride=1,padding=1)
        self.pool4=nn.MaxPool2d(kernel_size=2, stride=2)
        # 16*16
        self.conv5=nn.Conv2d(in_channels=48, out_channels=48,kernel_size=3,stride=1,padding=1)
        self.pool5=nn.MaxPool2d(kernel_size=2, stride=2)
        # 8*8
        self.conv6=nn.Conv2d(in_channels=48, out_channels=48,kernel_size=3,stride=1,padding=1)
        self.pool6=nn.MaxPool2d(kernel_size=2, stride=2)
        # 4*4
        self.conv7=nn.Conv2d(in_channels=48, out_channels=48,kernel_size=3,stride=1,padding=1)
        self.pool7=nn.MaxPool2d(kernel_size=2, stride=2)
        # 2*2
        
        # 定义多个专家（分类头）
        num_experts=10
        self.experts = nn.ModuleList([     # 这种方式避免共享权重
            nn.Sequential(
                nn.Linear(in_features=2*2*48, out_features=256),
                nn.ReLU(),
                nn.Linear(in_features=256, out_features=10)
            ) for _ in range(num_experts)
        ])
        # 定义门控单元
        self.gate=nn.Linear(in_features=2*2*48, out_features=10)
        
        self.activate=nn.ReLU()
        
        
    def forward(self,x, count=False):
        c1=self.conv1(x)
        p1=self.pool1(c1)
        a1=self.activate(p1)
        
        c2=self.conv2(a1)
        p2=self.pool2(c2)
        a2=self.activate(p2)
        
        c3=self.conv3(a2)
        p3=self.pool3(c3)
        a3=self.activate(p3)
        
        c4=self.conv4(a3)
        p4=self.pool4(c4)
        a4=self.activate(p4)
        
        c5=self.conv5(a4)
        p5=self.pool5(c5)
        a5=self.activate(p5)
        
        c6=self.conv6(a5)
        p6=self.pool6(c6)
        a6=self.activate(p6)
        
        c7=self.conv7(a6)
        p7=self.pool7(c7)
        a7=self.activate(p7)
        
        a7_flat = a7.view(a7.size(0), -1)  # 将特征图展平为一维向量
        
        # 分类头相关处理
        gate_weights=self.gate(a7_flat)
        max_weights, max_indices = torch.max(gate_weights, dim=1)
        
        #以下部分应当并行，可以并行
        results=[]
        for data, indice in zip(a7_flat, max_indices):
            if count:
                head_counter[indice]+=1
            results.append(self.experts[indice](data))
        res=torch.stack(results,dim=0)
        
        return res